from sympy.matrices.dense import Matrix, MatrixBase
from sympy.tensor.functions import shape
from sympy.core.singleton import S

# from sympy.stats import ChiSquared, density, E, variance
# from sympy.core.symbol import _symbol
from sympy.functions.elementary.miscellaneous import sqrt
from sympy.functions.special.error_functions import erf, erfinv


class ChiSquaredTest:
    """
    主要是对高中数学里的卡方检测进行抽象处理，便于以后计算处理。
    其实我现在越看这个卡方独立性检验越觉得里面有问题，但是不管怎样，我觉得还是先以完成任务为第一要务，不对其基本的理论基础做评价。
    """

    def __init__(self, m: MatrixBase) -> None:
        self.m = m

    @property
    def chi_squared_value(self):
        if self.m.shape == (2, 2):
            n = sum(self.m)
            a = self.m[0]
            b = self.m[1]
            c = self.m[2]
            d = self.m[3]
            return n * (a * d - b * c) ** 2 / ((a + b) * (a + c) * (c + d) * (b + d))

    @property
    def error_probability(self):
        return error_probability_of_chi_squared(self.chi_squared_value)

    @property
    def expand_matrix(self):
        m_shape = shape(self.m)
        M = self.m.col_insert(
            m_shape[-1], Matrix([sum(self.m.row(i)) for i in range(m_shape[0])])
        )
        return M.row_insert(
            m_shape[0], Matrix([[sum(M.col(i)) for i in range(m_shape[-1] + 1)]])
        )


def chi_squared_of_error_probability(p):
    y = 1 - p
    return 2 * (erfinv(y)) ** 2


def error_probability_of_chi_squared(chi_squared):
    return 1 - erf(sqrt(chi_squared / S(2)))
